from BinaryTree import BinaryTree
from lib.myseq import MySeq
from lib.align import SubstMatrix
from lib.mult_align import PairwiseAlignment, MultipleAlignment
from UPGMA import UPGMA
from HierarchicalClustering import HierarchicalClustering
from NumMatrix import NumMatrix

class BinaryTree_ex(BinaryTree):
        
    def size(self):
        if self.value >= 0:
            return (0,1)

        num_leafes = 0
        num_nodes = 0
        if (self.left != None):
            left_size = self.left.size()
        else:
            left_size = (0,0)
        if (self.right != None):  
            right_size = self.right.size() 
        else:
            right_size = (0,0)
        
        num_nodes += (left_size[0] + right_size[0] + 1)
        num_leafes += (left_size[1] + right_size[1])
        return num_nodes, num_leafes

    def exists_leaf(self, num):
        if self.value >= 0:
            if self.value == num:
                return True
            else:
                return False
        else:
            if self.left != None:
                if self.left.exists_leaf(num):
                    return True
            if self.right != None:
                if self.right.exists_leaf(num):
                    return True
        return False
    
    def common_ancestor(self, leaf1, leaf2):
        if self.value >= 0:
            return None
        if self.left.exists_leaf(leaf1):
            if self.left.exists_leaf(leaf2):
                return self.left.common_ancestor(leaf1,leaf2)
            elif self.right.exists_leaf(leaf2):
                return self
        elif self.right.exists_leaf(leaf1):
            if self.right.exists_leaf(leaf2):
                return self.right.common_ancestor(leaf1,leaf2)
            elif self.left.exists_leaf(leaf2):
                return self
        return None

    def common_ancestor_advance(self,leafs):
        if self.value >= 0:
            return None
        
        in_left = [self.left.exists_leaf(leaf) for leaf in leafs]
        in_right = [self.right.exists_leaf(leaf) for leaf in leafs]
        if all(in_left):
            return self.left.common_ancestor_advance(leafs)
        elif all(in_right):
            return self.right.common_ancestor_advance(leafs)
        else:
            in_tree = [l or r for l,r in zip(in_left,in_right)]
            if all(in_tree):
                return self
            else:
                return None

    def distance_leaves(self, leafnum1, leafnum2):
        ancestor = self.common_ancestor(leafnum1, leafnum2)
        return 2*ancestor.distance 
    
    def nearest_leaves(self):
        if self.distance == 0:
            raise ValueError('less than two leaves')
        elif self.right.distance == 0 and self.left.distance == 0:
            return self.distance * 2
        elif self.right.distance == 0:
            return self.left.nearest_leaves()
        elif self.left.distance == 0:
            return self.right.nearest_leaves()
        else:
            return min(self.left.nearest_leaves(),self.right.nearest_leaves())

def UPGMA_case1():
    s1 = MySeq("ACATATCAT")
    s2 = MySeq("AGATATTAG")
    s3 = MySeq("AACAGATCT")
    s4 = MySeq("GCATCGATT")
    
    sm = SubstMatrix()
    sm.create_submat(1,-1,"ACGT")
    aligner = PairwiseAlignment(sm,-1)
    aligenment = MultipleAlignment([s1, s2, s3, s4], aligner).align_consensus()
    print(aligenment)

    up  = UPGMA([s1, s2, s3, s4], aligner)
    arv = up.run()    
    arv.print_tree() 

def create_tree():
    mat = NumMatrix(4,4)
    mat.set_value(1,0,8)
    mat.set_value(2,0,8)
    mat.set_value(3,0,8)
    mat.set_value(2,1,6)
    mat.set_value(3,2,6)
    mat.set_value(3,1,4)
    model = HierarchicalClustering(mat)
    tree = model.execute_clustering()
    tree.print_tree()

def test_advance_BinaryTree():              
    a = BinaryTree_ex(1)
    b = BinaryTree_ex(2)
    c = BinaryTree_ex(3)
    d = BinaryTree_ex(4)
    e = BinaryTree_ex(-1, 2.0, b, c)
    f = BinaryTree_ex(-1, 1.5, d, a)
    g = BinaryTree_ex(-1, 4.5, e, f)
    g.print_tree()

    print("a:")
    print(g.size())

    print('b:')
    print(g.exists_leaf(1))
    print(g.exists_leaf(5))

    print('c:')
    g.common_ancestor(1,4).print_tree()
    g.common_ancestor(2,3).print_tree()
    g.common_ancestor(2,4).print_tree()

    print('d:')
    g.common_ancestor_advance([1,4]).print_tree()
    h = BinaryTree_ex(5)
    i = BinaryTree_ex(-1, 6, g, h)
    i.common_ancestor_advance([1,2,4]).print_tree()
    i.common_ancestor_advance([4,5]).print_tree()

    print('e:')
    print(i.distance_leaves(1,4))
    print(i.distance_leaves(1,2))

    print('f:')
    print(i.nearest_leaves())



if __name__ == "__main__":
    print("### ex1 ###")
    UPGMA_case1()
    print("### ex2 ###")
    create_tree()
    print("### ex3 ###")
    test_advance_BinaryTree()
